import sqlite3
import pandas as pd
from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from datetime import datetime, timedelta, time, date
import pytz
import os

# 定义旧数据库和新数据库的路径
OLD_DB_PATH = 'db/trades.db.bak'
NEW_DB_PATH = 'db/trades.db'

# 定义 SQLAlchemy 模型 (与您 src/core/models.py 中的定义一致，但只用于迁移)
Base = declarative_base()

class Trade(Base):
    __tablename__ = 'trades'
    index = Column(Integer, primary_key=True)
    trade_id = Column(String, unique=True)
    symbol = Column(String)
    open_datetime = Column(DateTime)
    open_price = Column(Float)
    open_qty = Column(Float)
    close_datetime = Column(DateTime)
    close_price = Column(Float)
    close_qty = Column(Float)
    direction = Column(String)
    open_order_type = Column(String)
    trade_type = Column(String)
    side = Column(String)
    point_profit_per_lot = Column(Float)
    total_points = Column(Float)
    gross_profit = Column(Float)
    net_profit = Column(Float)
    commission = Column(Float)
    holding_time = Column(String)
    nth_trade = Column(Integer)
    session = Column(String)
    max_floating_profit_points = Column(Float)
    max_floating_profit_dollars = Column(Float)
    max_floating_loss_points = Column(Float)
    max_floating_loss_dollars = Column(Float)
    entry_reason = Column(String)
    exit_reason = Column(String)
    stop_loss_price = Column(Float)
    stop_loss_reason = Column(String)
    take_profit_price = Column(Float)
    take_profit_reason = Column(String)

class Order(Base):
    __tablename__ = 'orders'
    index = Column(Integer, primary_key=True)
    order_id = Column(String, unique=True)
    symbol = Column(String)
    side = Column(String)
    order_type = Column(String)
    quantity = Column(Float)
    filled = Column(Float)
    quantity_left = Column(Float)
    limit_price = Column(Float)
    stop_price = Column(Float)
    avg_fill_price = Column(Float)
    status = Column(String)
    open_time = Column(DateTime)
    close_time = Column(DateTime)
    duration = Column(String)
    commission = Column(Float)
    expiration = Column(DateTime)
    multiplier = Column(Float)

# 辅助函数：根据纽约时间判断纽约证券交易所的交易日
def get_nyse_trading_day(dt: datetime) -> date:
    if dt.time() >= time(9, 30):
        return dt.date()
    else:
        return (dt - timedelta(days=1)).date()

def safe_float(value):
    try:
        return float(value)
    except (ValueError, TypeError):
        return None

def safe_datetime(value):
    if isinstance(value, datetime):
        return value
    if isinstance(value, str):
        try:
            return datetime.fromisoformat(value)
        except ValueError:
            pass
        for fmt in ('%Y-%m-%d %H:%M:%S.%f', '%Y-%m-%d %H:%M:%S', '%Y-%m-%d %H:%M', '%Y/%m/%d %H:%M'):
            try:
                return datetime.strptime(value, fmt)
            except ValueError:
                continue
    return None

def migrate_trades():
    print(f"旧数据库路径: {OLD_DB_PATH}")
    print(f"新数据库路径: {NEW_DB_PATH}")
    print("开始迁移 trades 表...")
    # 连接到旧数据库并读取数据
    try:
        old_conn = sqlite3.connect(OLD_DB_PATH)
        print(f"成功连接到旧数据库: {OLD_DB_PATH}")
    except Exception as e:
        print(f"连接旧数据库 {OLD_DB_PATH} 失败: {e}")
        return
    
    try:
        old_df = pd.read_sql_query("SELECT * FROM trades", old_conn)
        print(f"从旧数据库读取 trades 表数据成功。")
    except Exception as e:
        print(f"读取 trades 表数据失败: {e}")
        old_conn.close()
        return
    old_conn.close()

    if old_df.empty:
        print("旧数据库中没有 trades 表数据，跳过 trades 迁移。")
        return

    print(f"从旧数据库读取到 {len(old_df)} 条交易记录。")

    # 处理日期时间列，确保它们是 datetime 对象
    old_df['open_datetime'] = old_df['open_datetime'].apply(safe_datetime)
    old_df['close_datetime'] = old_df['close_datetime'].apply(safe_datetime)

    # 重新计算 nth_trade
    # 假设数据库中的 open_datetime 是北京时间（naive），先本地化再转换为纽约时区
    beijing_tz = pytz.timezone('Asia/Shanghai')
    new_york_tz = pytz.timezone('America/New_York')

    trades_df_for_nth = old_df.copy() # 使用副本进行计算
    if 'open_datetime' in trades_df_for_nth.columns:
        trades_df_for_nth['open_datetime'] = pd.to_datetime(trades_df_for_nth['open_datetime'], errors='coerce')
        trades_df_for_nth = trades_df_for_nth.dropna(subset=['open_datetime']) # 过滤掉无法解析的日期时间

        if not trades_df_for_nth.empty:
            trades_df_for_nth['open_datetime_ny'] = trades_df_for_nth['open_datetime'].dt.tz_localize(beijing_tz).dt.tz_convert(new_york_tz)
            trades_df_for_nth['ny_trading_day'] = trades_df_for_nth['open_datetime_ny'].apply(get_nyse_trading_day)
            
            trades_df_for_nth = trades_df_for_nth.sort_values(by=['ny_trading_day', 'open_datetime'])
            trades_df_for_nth['nth_trade_calculated'] = trades_df_for_nth.groupby('ny_trading_day').cumcount() + 1
            
            # 将计算出的 nth_trade 更新回原始 old_df
            # 使用 merge 或 map 来更新，避免循环
            old_df = old_df.set_index('index')
            trades_df_for_nth = trades_df_for_nth.set_index('index')
            old_df['nth_trade'] = trades_df_for_nth['nth_trade_calculated']
            old_df = old_df.reset_index() # 恢复index列
            # 确保 nth_trade 是整数类型
            old_df['nth_trade'] = old_df['nth_trade'].fillna(0).astype(int)
            print("nth_trade 重新计算完成。")

    # 连接到新数据库
    try:
        new_engine = create_engine(f'sqlite:///{NEW_DB_PATH}')
        Base.metadata.create_all(new_engine) # 确保表结构存在
        print(f"成功连接到新数据库并确保表结构存在: {NEW_DB_PATH}")
    except Exception as e:
        print(f"连接新数据库 {NEW_DB_PATH} 失败或创建表结构失败: {e}")
        return
    
    Session = sessionmaker(bind=new_engine)
    new_session = Session()
    print("新数据库会话已创建。")

    # 映射 DataFrame 列名到模型属性
    trade_columns_mapping = {
        col.name: col.name for col in Trade.__table__.columns
    }
    
    # 将 DataFrame 数据转换为 ORM 对象并保存
    migrated_count = 0
    for index, row in old_df.iterrows():
        try:
            trade_data = {k: row[k] for k in trade_columns_mapping.keys() if k in row}
            new_trade = Trade(**trade_data)
            new_session.add(new_trade)
            migrated_count += 1
            if migrated_count % 100 == 0:
                print(f"已处理 {migrated_count} 条交易记录...")
        except Exception as e:
            print(f"迁移交易记录 {row.get('trade_id', row.get('index', '未知'))} 时出错: {e}")
            new_session.rollback() # 回滚当前事务，尝试继续下一个
            continue
    
    try:
        new_session.commit()
        print(f"成功将 {migrated_count} 条交易记录提交到新数据库。")
    except Exception as e:
        print(f"提交交易记录到新数据库失败: {e}")
        new_session.rollback() # 再次回滚以防万一
    finally:
        new_session.close()
        print("新数据库会话已关闭。")

def migrate_orders():
    print("开始迁移 orders 表...")
    # 连接到旧数据库并读取数据
    try:
        old_conn = sqlite3.connect(OLD_DB_PATH)
        print(f"成功连接到旧数据库: {OLD_DB_PATH}")
    except Exception as e:
        print(f"连接旧数据库 {OLD_DB_PATH} 失败: {e}")
        return

    try:
        old_df = pd.read_sql_query("SELECT * FROM orders", old_conn)
        print(f"从旧数据库读取 orders 表数据成功。")
    except Exception as e:
        print(f"读取 orders 表数据失败: {e}")
        old_conn.close()
        return
    old_conn.close()

    if old_df.empty:
        print("旧数据库中没有 orders 表数据，跳过 orders 迁移。")
        return
        
    print(f"从旧数据库读取到 {len(old_df)} 条订单记录。")

    # 处理日期时间列，确保它们是 datetime 对象
    old_df['open_time'] = old_df['open_time'].apply(safe_datetime)
    old_df['close_time'] = old_df['close_time'].apply(safe_datetime)
    old_df['expiration'] = old_df['expiration'].apply(safe_datetime)

    # 连接到新数据库
    try:
        new_engine = create_engine(f'sqlite:///{NEW_DB_PATH}')
        Base.metadata.create_all(new_engine) # 确保表结构存在
        print(f"成功连接到新数据库并确保表结构存在: {NEW_DB_PATH}")
    except Exception as e:
        print(f"连接新数据库 {NEW_DB_PATH} 失败或创建表结构失败: {e}")
        return
    
    Session = sessionmaker(bind=new_engine)
    new_session = Session()
    print("新数据库会话已创建。")

    # 映射 DataFrame 列名到模型属性
    order_columns_mapping = {
        col.name: col.name for col in Order.__table__.columns
    }

    # 将 DataFrame 数据转换为 ORM 对象并保存
    migrated_count = 0
    for index, row in old_df.iterrows():
        try:
            order_data = {k: row[k] for k in order_columns_mapping.keys() if k in row}
            new_order = Order(**order_data)
            new_session.add(new_order)
            migrated_count += 1
            if migrated_count % 100 == 0:
                print(f"已处理 {migrated_count} 条订单记录...")
        except Exception as e:
            print(f"迁移订单记录 {row.get('order_id', row.get('index', '未知'))} 时出错: {e}")
            new_session.rollback() # 回滚当前事务，尝试继续下一个
            continue

    try:
        new_session.commit()
        print(f"成功将 {migrated_count} 条订单记录提交到新数据库。")
    except Exception as e:
        print(f"提交订单记录到新数据库失败: {e}")
        new_session.rollback()
    finally:
        new_session.close()
        print("新数据库会话已关闭。")

if __name__ == "__main__":
    # 确保 db 目录存在
    os.makedirs(os.path.dirname(NEW_DB_PATH), exist_ok=True)
    print(f"确保目录 {os.path.dirname(NEW_DB_PATH)} 存在。")
    
    # 确保新数据库文件存在 (如果不存在，SQLAlchemy会自动创建)
    # 这一步不是严格必需的，但可以明确提示文件路径
    if not os.path.exists(NEW_DB_PATH):
        print(f"创建新的数据库文件: {NEW_DB_PATH}")
        # 如果文件不存在，可以手动创建空文件，或者让SQLAlchemy在第一次连接时创建
        # with open(NEW_DB_PATH, 'a'):
        #    os.utime(NEW_DB_PATH, None)
    
    # 执行迁移
    migrate_trades()
    migrate_orders()
    print("数据迁移完成！") 